package com.itcast.stargraph.service.impl;

import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONUtil;
import com.alibaba.fastjson2.JSON;
import com.alibaba.fastjson2.JSONObject;
import com.itcast.stargraph.constant.CacheConstant;
import com.itcast.stargraph.dto.request.Text2ImageReqDto;
import com.itcast.stargraph.dto.respone.Text2ImageResDto;
import com.itcast.stargraph.pojo.ComfyuiModel;
import com.itcast.stargraph.pojo.ComfyuiRequestDto;
import com.itcast.stargraph.pojo.ComfyuiTask;
import com.itcast.stargraph.service.FreemarkerService;
import com.itcast.stargraph.service.TextToImageService;
import com.itcast.stargraph.utils.OllamaUtils;
import com.itcast.stargraph.utils.UserThreadLocalUtils;
import org.springframework.beans.BeanUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.ZSetOperations;
import org.springframework.stereotype.Service;

import java.time.Duration;

/**
 * 文生图服务实现类
 *
 * @Author katana
 * @Date 2025/6/11 14:19
 */
@Service
public class TextToImageServiceImpl implements TextToImageService {
    
    @Autowired
    private OllamaUtils ollamaUtils;
    
    @Autowired
    private FreemarkerService freemarkerService;
    
    @Autowired
    private StringRedisTemplate redisTemplate;
    
    
    /**
     * 获取ComfyuiTask任务实例
     *
     * @param dto
     * @return
     */
    @Override
    public ComfyuiTask getComfyuiTask(Text2ImageReqDto dto) {
        //1 将dto转换成ComfyuiModel
        ComfyuiModel model = new ComfyuiModel();
        BeanUtils.copyProperties(dto, model);
        
        //2 补充model的属性
        //2.1 将请求提示翻译成中文
        String prompt = ollamaUtils.translate(dto.getPrompt());
        if (!dto.getReverse().isBlank()){
            String resverse = ollamaUtils.translate(dto.getReverse());
            model.setReverse(resverse);
        }
        
        
        model.setPrompt(prompt);
        model.setModelName(dto.modelName());
        model.setWidth(dto.width());
        model.setHeight(dto.height());
        model.setScheduler(dto.scheduler());
        model.setSamplerName(dto.samplerName());
        
        // 3 渲染工作流json
        String workFlowJson = freemarkerService.renderWorkFlowJSON(model);
        // 将字符串解析为JSONObject
        JSONObject promptObj = JSON.parseObject(workFlowJson);
        // 4 封装ComfyuiTask对象
        //这里设置WS_CLIENT_ID,是comfyui服务端返回的clientId
        ComfyuiRequestDto comfyuiRequestDto = new ComfyuiRequestDto(CacheConstant.WS_CLIENT_ID, promptObj);
        //这里设置的clientId是前端的clientId
        ComfyuiTask comfyuiTask = new ComfyuiTask(dto.getClientId(), UserThreadLocalUtils.getUser().getId(), comfyuiRequestDto);
        return comfyuiTask;
    }
    
    /**
     * 添加ComfyuiTask任务实例到队列中排队
     *
     * @param task
     * @return
     */
    @Override
    public ComfyuiTask addQueueTask(ComfyuiTask task) {
        //1.基于redis计数器+1，获取最新计数结果，当做任务存储到队列中的score
        Long score = redisTemplate.opsForValue().increment(CacheConstant.REDIS_KEY_COUNTER);
        task.setIndex(score);
        
        //2.添加任务到redis的zset队列中排队（value：队列ID，score：自增ID值）
        redisTemplate.opsForZSet().add(CacheConstant.REDIS_KEY_QUEUE, task.getId(), score);
        
        //3.添加任务的JSON到redis的string中
        redisTemplate.opsForValue().set(String.format(CacheConstant.REDIS_KEY_TASK_PREFIX,task.getId()), JSONUtil.toJsonStr(task));
        
        return task;
    }
    
    /**
     * 添加comfyui流程任务（工作流任务）
     *
     * @param dto
     * @return
     */
    @Override
    public Text2ImageResDto prompt(Text2ImageReqDto dto) {
        //1. 获取comfyui任务实例
        ComfyuiTask comfyuiTask = this.getComfyuiTask(dto);
        
        //2. 添加任务到redis队列中排队
        ComfyuiTask task = this.addQueueTask(comfyuiTask);
        
        //3. 封装响应实例
        Text2ImageResDto resDto = new Text2ImageResDto();
        //任务在队列中的排序值
        resDto.setQueueIndex(task.getIndex());
        //任务id
        resDto.setPid(task.getId());
        return resDto;
    }
    
    
    /**
     * 从任务队列中获取ComfyuiTask任务实例
     *
     * @return
     */
    @Override
    public ComfyuiTask popQueueTask() {
        //1 获取任务队列,并判断队列是否为空
        Long size = redisTemplate.opsForZSet().size(CacheConstant.REDIS_KEY_QUEUE);
        if (size == null || size == 0) {
            return null;
        }
        //2 获取队列中值最小的任务(优先级最高)
        ZSetOperations.TypedTuple<String> task = redisTemplate.opsForZSet().popMin(CacheConstant.REDIS_KEY_QUEUE);
        if (task == null){
            return null;
        }
        String taskId = task.getValue();
        
        //3 根据任务ID获取任务实例
        String taskJson = redisTemplate.opsForValue().get(String.format(CacheConstant.REDIS_KEY_TASK_PREFIX, taskId));
        if (taskJson == null){
            return null;
        }
        ComfyuiTask comfyuiTask = JSONUtil.toBean(taskJson, ComfyuiTask.class);
        
        return comfyuiTask;
    }
    
    
    /**
     * 缓存最开始的任务
     *
     * @param prompId 任务ID
     * @param task    最先开始的任务
     */
    @Override
    public void addStartedTask(String prompId, ComfyuiTask task) {
        //更新任务缓存
        redisTemplate.opsForValue().set(String.format(CacheConstant.REDIS_KEY_TASK_PREFIX, prompId), JSONUtil.toJsonStr(task), Duration.ofMinutes(30));
        
        //redisTemplate.opsForValue().set(String.format(CacheConstant.WS_TASK_REDIS_KEY,prompId),JSONUtil.toJsonStr(task), Duration.ofMinutes(30));
    }
    
    /**
     * 从缓存取出最开始的任务
     *
     * @param prompId 任务ID
     */
    @Override
    public ComfyuiTask getStartedTask(String prompId) {
        //String taskJSON = redisTemplate.opsForValue().get(String.format(CacheConstant.WS_TASK_REDIS_KEY, prompId));
        String taskJSON = redisTemplate.opsForValue().get(String.format(CacheConstant.REDIS_KEY_TASK_PREFIX, prompId));
        if(!StrUtil.isBlank(taskJSON)) {
            return JSONUtil.toBean(taskJSON,ComfyuiTask.class);
        }
        return null;
    }
}
